import os
import pandas as pd
import cv2
from pathlib import Path
import glob


def get_crop_locs(edge_len, base_size):
    edge_times =  edge_len // base_size
    edge_crops = []
    edge_locs = []

    if edge_times == 0:
        edge_times == 1
        edge_crops = [edge_len]
        edge_locs = [0]
    else:
        carry = edge_len % base_size
        
        if carry >= 100:
            edge_locs = [i*base_size for i in range(edge_times)]
            edge_locs.append(edge_len-base_size)
            edge_crops = [base_size] * (edge_times + 1)
        else:
            base_size_dilate =  base_size + carry // edge_times
            edge_locs = [i*base_size_dilate for i in range(edge_times-1)]
            last_base_size = base_size_dilate + carry % edge_times
            edge_locs.append((edge_times-1) * base_size_dilate)
            edge_crops = [base_size_dilate] * (edge_times -1) + [last_base_size]
    return edge_locs, edge_crops

def adaptive_crop(img, base_size=1024):
    h, w = img.shape[:2]

    ys, h_crops = get_crop_locs(h, base_size)
    xs, w_crops = get_crop_locs(w, base_size)
    return xs, ys, h_crops, w_crops

def img_split(img_path, mask_root, img_dst_root, mask_dst_root):
    img = cv2.imread(img_path, 1)
    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
    code = Path(img_path).parent.name
    name_stem = Path(img_path).stem
    mask_path = os.path.join(mask_root, code, name_stem+'.png')
    if not os.path.exists(mask_path):
        return

    mask = cv2.imread(mask_path, 0)

    xs, ys, h_crops, w_crops = adaptive_crop(img)

    index = 1
    for x, w_crop in zip(xs, w_crops):
        for y, h_crop in zip(ys, h_crops):
            cur_img = img[y:y+h_crop, x:x+w_crop]
            cur_mask = mask[y:y+h_crop, x:x+w_crop]

            if 255 in cur_mask:
                cur_code = code
                os.makedirs(os.path.join(img_dst_root, cur_code), exist_ok=True)
                cv2.imwrite(os.path.join(img_dst_root, cur_code, name_stem+"_"+str(index)+'.jpg'), cur_img)
                os.makedirs(os.path.join(mask_dst_root, cur_code), exist_ok=True)
                cv2.imwrite(os.path.join(mask_dst_root, cur_code, name_stem+"_"+str(index)+'.png'), cur_mask)
            else:
                cur_code = 'TGXID'
                os.makedirs(os.path.join(img_dst_root, cur_code), exist_ok=True)
                cv2.imwrite(os.path.join(img_dst_root, cur_code, name_stem+"_"+str(index)+'.jpg'), cur_img)
            
            index += 1
            

def main():
    try:
        from pandarallel import pandarallel
        pandarallel.initialize(progress_bar=True) 
        print('Use multi threading !')
        is_pandarallel = True
    except:
        print('Use single threading !')
        is_pandarallel = False

    img_root = r'/data2/autorepair/ruanzhifeng/autorepair_t7_10/t7/T6007/all_0409'
    mask_root = r'/data2/autorepair/ruanzhifeng/autorepair_t7_10/t7/T6007/all_0409_mask'

    img_dst_root = r'/data2/autorepair/ruanzhifeng/autorepair_t7_10/t7/T6007/cropped_img_0409'
    mask_dst_root = r'/data2/autorepair/ruanzhifeng/autorepair_t7_10/t7/T6007/cropped_mask_0409'

    df = pd.DataFrame()
    df['img_path'] = glob.glob(os.path.join(img_root, "*/*.jpg"))
    print(len(df))

    if is_pandarallel:
        df.parallel_apply(lambda x:img_split(img_path=x['img_path'], mask_root=mask_root, img_dst_root=img_dst_root, mask_dst_root=mask_dst_root), axis=1)
    else:
        df.apply(lambda x:img_split(img_path=x['img_path'], mask_root=mask_root, img_dst_root=img_dst_root, mask_dst_root=mask_dst_root), axis=1)



if __name__=='__main__':
    main()
